import java.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map.Entry;

/*
 * TNBClassifier.java
 * Usage: IR HW4. This class is subclass of AbstractClassifier. it overrides the train() method using first 3 step introduce in the paper
 */
public class TNBClassifier extends AbstractClassifier {

	// private double[][] w;
	private int termNum;

	public TNBClassifier() {
		super();
	}

	/**
	 * @param args
	 */
	public static void main(String[] args) {

		AbstractClassifier classifier = new TNBClassifier();
		classifier.train("citeseer.train.txt");
		classifier.classify("citeseer.test.txt");
		classifier.evaluate();
	}

	/*
	 * it overrides the train() method using first 3 step introduce in the paper
	 */
	public void train(String file) {
		try {
			// total number of training documents
			int docNum = 0;
			// number of documents in each category
			HashMap<Integer, Integer> docNumInClass = new HashMap<Integer, Integer>();
			// terms appearing in all categories
			HashSet<Integer> terms = new HashSet<Integer>();
			// term frequency in each document
			HashMap<Integer, HashMap<Integer, Double>> termInDoc = new HashMap<Integer, HashMap<Integer, Double>>();
			// training labels of each document
			HashMap<Integer, Integer> docClass = new HashMap<Integer, Integer>();
			// all categories
			HashSet<Integer> totalClasses = new HashSet<Integer>();

			BufferedReader br = new BufferedReader(new FileReader(file));
			String line;
			while ((line = br.readLine()) != null) {

				String[] str = line.split(" ");
				String[] classes = str[0].split(",");
				for (int i = 0; i < classes.length; i++) {
					int c = Integer.parseInt(classes[i]);

					if (!totalClasses.contains(c)) {
						totalClasses.add(c);
					}

					if (docNumInClass.containsKey(c)) {
						docNumInClass.put(c, docNumInClass.get(c) + 1);
					} else {
						docNumInClass.put(c, 1);
					}

					HashMap<Integer, Double> termFreq = new HashMap<Integer, Double>();
					docNum++;
					termInDoc.put(docNum, termFreq);
					docClass.put(docNum, c);

					for (int j = 1; j < str.length - 2; j++) {
						// System.out.println(str[i]);
						String[] pair = str[j].split(":");
						int term = Integer.parseInt(pair[0]);
						double times = Double.parseDouble(pair[1]);

						if (!terms.contains(term)) {
							terms.add(term);
						}

						termFreq.put(term, times);

					}
				}
			}

			br.close();

			for (Iterator<Entry<Integer, Integer>> it = docNumInClass
					.entrySet().iterator(); it.hasNext();) {
				Entry<Integer, Integer> e = it.next();
				pc.put(e.getKey(), (double) e.getValue() / (double) docNum);
			}

			//System.out.println(pc);

			// System.out.println(termInDoc);
			// System.out.println(docClass);
			termNum = terms.size();

			/*
			 * get matix d
			 */
			double[][] d = new double[docNum + 1][termNum + 1];

			for (int j = 1; j <= docNum; j++) {
				HashMap<Integer, Double> termFreq = termInDoc.get(j);
				for (int i = 1; i <= termNum; i++) {
					if (termFreq.containsKey(i)) {
						d[j][i] = termFreq.get(i);
					} else {
						d[j][i] = 0.0;
					}
				}
			}

			//print(d);

			// TWCNB
			// step 1
			for (int j = 1; j <= docNum; j++) {
				for (int i = 1; i <= termNum; i++) {

					d[j][i] = Math.log(d[j][i] + 1.0);
				}
			}

			//print(d);
			// step2

			for (int i = 1; i <= termNum; i++) {
				int occurrence = 0;
				for (int j = 1; j <= docNum; j++) {
					if (d[j][i] != 0.0) {
						occurrence++;
					}
				}
				// System.out.println(occurrence);

				double idf = Math.log((double) docNum / (double) occurrence);
				// System.out.println(idf);

				// System.out.println(idf);
				for (int j = 1; j <= docNum; j++) {
					d[j][i] = d[j][i] * idf;
				}
			}
			//print(d);
//
			// step3
			for (int j = 1; j <= docNum; j++) {
				double length = 0;
				for (int i = 1; i <= termNum; i++) {
					length += d[j][i] * d[j][i];
				}

				length = Math.sqrt(length);

				for (int i = 1; i <= termNum; i++) {
					d[j][i] = d[j][i] / length;
				}
			}

			//print(d);

			classNum = totalClasses.size();

			for (int c = 1; c <= classNum; c++) {
				// System.out.println("class=" + c);
				double totalNumInClass = 0;
				double[] termNumInClass = new double[termNum + 1];
				for (int j = 1; j <= docNum; j++) {
					if (docClass.get(j) == c) {
						// System.out.println(j);
						for (int i = 1; i <= termNum; i++) {
							totalNumInClass += d[j][i];
							termNumInClass[i] += d[j][i];
						}
					}
				}
				HashMap<Integer, Double> termInClass = new HashMap<Integer, Double>();
				for(int i = 1; i <=termNum; i++){
					termInClass.put(i, (termNumInClass[i] + 1) / (totalNumInClass + termNum));
				}
				ptc.put(c, termInClass);

			}
			//System.out.println(ptc);
		} catch (FileNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}

	}

	/*
	 * helper method to print matrix
	 */
	public void print(double[][] d) {
		for (int j = 1; j < d.length; j++) {
			for (int i = 1; i < d[j].length; i++) {
				System.out.print(d[j][i] + "\t\t");
			}
			System.out.println();
		}
		System.out.println("======");
	}

}
